import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

import sys
import os.path as osp
sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))
from sam_pth.common_tools import *


def main():
    device = device_select()
    sam2_checkpoint = "/weights/sam2.1_hiera_small.pt"
    model_cfg = "//opt/sam2/sam2/configs/sam2.1/sam2.1_hiera_s.yaml"
    predictor = init_predictor(sam2_checkpoint, model_cfg, device)

    image_path = '/opt/sam2/notebooks/images/truck.jpg'
    image2_path = '/opt/sam2/notebooks/images/groceries.jpg'
    save_dir = '/data'
    image = Image.open(image_path)
    image = np.array(image.convert("RGB"))

    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    plt.axis('on')
    if show_image:
        plt.show()
    else:
        plt.savefig(osp.join(save_dir, 'truck-1.jpg'))

    # Process the image to produce an image embedding by calling SAM2ImagePredictor.set_image.
    # SAM2ImagePredictor remembers this embedding and will use it for subsequent mask prediction.
    predictor.set_image(image)
    print(predictor._features["image_embed"].shape)

    # 设置prompt point和类别
    input_point = np.array([[500, 375]])
    input_label = np.array([1])

    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    show_points(input_point, input_label, plt.gca())
    plt.axis('on')
    if show_image:
        plt.show()
    else:
        plt.savefig(osp.join(save_dir, 'truck-2.jpg'))

    # 当multimask_output = True（默认设置）时，SAM2会输出3个掩码，其中scores给出了模型对这些掩码质量的自身评估。
    # 此设置适用于模糊的输入提示，有助于模型消除与提示相符的不同对象的歧义。
    # 当设置为False时，它将仅返回一个掩码。对于像单个点这样的模糊提示，即使只需要一个掩码，也建议使用multimask_output = True；
    # 可以通过选择scores中返回分数最高的那个掩码来选出最佳的单个掩码，这样通常会得到更好的掩码。
    masks, scores, logits = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        multimask_output=True,
    )

    sorted_ind = np.argsort(scores)[::-1]
    masks = masks[sorted_ind]
    scores = scores[sorted_ind]
    logits = logits[sorted_ind]

    show_masks(image, masks, scores, point_coords=input_point, input_labels=input_label, borders=True, save_path=osp.join(save_dir, 'truck---1.jpg'))

    # 单个输入点具有模糊性，模型会返回多个与之相符的对象。
    # 若要获取单个对象，可以提供多个点。如果有前一次迭代得到的掩码，也可以将其提供给模型辅助预测。
    # 当使用多个提示指定单个对象时，可通过将multimask_output设置为False来只获取一个掩码。
    input_point = np.array([[500, 375], [1125, 625]])
    input_label = np.array([1, 1])
    mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask

    masks, scores, _ = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        mask_input=mask_input[None, :, :],
        multimask_output=False,
    )

    show_masks(image, masks, scores, point_coords=input_point, input_labels=input_label, save_path=osp.join(save_dir, 'truck---2.jpg'))

    # 为了排除汽车，仅指定车窗，可以提供一个背景点（标签为0，此处显示为红色）。
    input_point = np.array([[500, 375], [1125, 625]])
    input_label = np.array([1, 0])
    mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask

    masks, scores, _ = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        mask_input=mask_input[None, :, :],
        multimask_output=False,
    )

    show_masks(image, masks, scores, point_coords=input_point, input_labels=input_label, save_path=osp.join(save_dir, 'truck---3.jpg'))

    # 该模型也可以接受以xyxy格式给出的边界框作为输入。
    input_box = np.array([425, 600, 700, 875])

    masks, scores, _ = predictor.predict(
        point_coords=None,
        point_labels=None,
        box=input_box[None, :],
        multimask_output=False,
    )

    show_masks(image, masks, scores, box_coords=input_box, save_path=osp.join(save_dir, 'truck---4.jpg'))

    # 点和边界框可以组合使用，只需将这两种类型的提示信息都提供给预测器即可。在此例中，这种方式可用于仅选择卡车轮胎，而非整个车轮。
    input_box = np.array([425, 600, 700, 875])
    input_point = np.array([[575, 750]])
    input_label = np.array([0])

    masks, scores, logits = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        box=input_box,
        multimask_output=False,
    )

    show_masks(image, masks, scores, box_coords=input_box, point_coords=input_point, input_labels=input_label, save_path=osp.join(save_dir, 'truck---5.jpg'))

    # SAM2图像预测器可以使用predict方法为同一幅图像接收多个输入提示。例如，假设我们从一个物体检测器得到了几个边界框输出。
    input_boxes = np.array([
        [75, 275, 1725, 850],
        [425, 600, 700, 875],
        [1375, 550, 1650, 800],
        [1240, 675, 1400, 750],
    ])

    masks, scores, _ = predictor.predict(
        point_coords=None,
        point_labels=None,
        box=input_boxes,
        multimask_output=False,
    )

    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    for mask in masks:
        show_mask(mask.squeeze(0), plt.gca(), random_color=True)
    for box in input_boxes:
        show_box(box, plt.gca())
    plt.axis('off')
    if show_image:
        plt.show()
    else:
        plt.savefig(osp.join(save_dir, 'truck-3.jpg'))

    # End - to - end batched inference
    image1 = image  # truck.jpg from above
    image1_boxes = input_boxes
    image2 = Image.open(image2_path)
    image2 = np.array(image2.convert("RGB"))
    image2_boxes = np.array([
        [450, 170, 520, 350],
        [350, 190, 450, 350],
        [500, 170, 580, 350],
        [580, 170, 640, 350],
    ])

    img_batch = [image1, image2]
    boxes_batch = [image1_boxes, image2_boxes]

    predictor.set_image_batch(img_batch)

    masks_batch, scores_batch, _ = predictor.predict_batch(
        None,
        None,
        box_batch=boxes_batch,
        multimask_output=False
    )

    for i, (image, boxes, masks) in enumerate(zip(img_batch, boxes_batch, masks_batch)):
        plt.figure(figsize=(10, 10))
        plt.imshow(image)
        for mask in masks:
            show_mask(mask.squeeze(0), plt.gca(), random_color=True)
        for box in boxes:
            show_box(box, plt.gca())
        plt.axis('off')
        if show_image:
            plt.show()
        else:
            plt.savefig(osp.join(save_dir, 'zip0-batch-' + str(i) + '.jpg'))

    image1_pts = np.array([
        [[500, 375]],
        [[650, 750]]
    ])  # Bx1x2 where B corresponds to number of objects
    image1_labels = np.array([[1], [1]])
    image2_pts = np.array([
        [[400, 300]],
        [[630, 300]],
    ])
    image2_labels = np.array([[1], [1]])

    pts_batch = [image1_pts, image2_pts]
    labels_batch = [image1_labels, image2_labels]

    masks_batch, scores_batch, _ = predictor.predict_batch(pts_batch,
                                                           labels_batch,
                                                           box_batch=None,
                                                           multimask_output=True)

    # Select the best single mask per object
    best_masks = []
    for masks, scores in zip(masks_batch, scores_batch):
        best_masks.append(masks[range(len(masks)), np.argmax(scores, axis=-1)])

    for i, (image, points, labels, masks) in enumerate(zip(img_batch, pts_batch, labels_batch, best_masks)):
        plt.figure(figsize=(10, 10))
        plt.imshow(image)
        for mask in masks:
            show_mask(mask, plt.gca(), random_color=True)
        show_points(points, labels, plt.gca())
        plt.axis('off')
        if show_image:
            plt.show()
        else:
            plt.savefig(osp.join(save_dir, 'zip1-batch-' + str(i) + '.jpg'))


if __name__ == '__main__':
    main()
